Skip to content

fix(ai-proxy): account for Bedrock prompt cache tokens - #14044

Open
AlinsRan wants to merge 1 commit into
apache:masterfrom
AlinsRan:fix/bedrock-cache-usage
Open

AlinsRan wants to merge 1 commit into
apache:masterfrom
AlinsRan:fix/bedrock-cache-usage

Conversation

@AlinsRan

@AlinsRan AlinsRan commented Oct 15, 2026 •

Copy link
Copy Markdown
Contributor

Description

Bedrock Converse reports cache-read and cache-write tokens separately from inputTokens. APISIX currently omits these tokens from normalized prompt usage and reports zero for both cache counters, affecting token accounting and LLM logs in streaming and non-streaming responses.

For example, a response with 10 input tokens, 20 cache-read tokens, 30 cache-write tokens, and 8 output tokens should report 60 prompt tokens, rather than 10.

Changes

Use the same usage normalization for ordinary responses and streaming metadata:

  • Include cache-read and cache-write tokens in prompt_tokens, and expose both cache counters to the existing LLM logging variables.
  • Preserve upstream totalTokens when present; otherwise, calculate the total from normalized prompt and completion tokens.
  • Accept the cacheReadInputTokenCount and cacheWriteInputTokenCount aliases. Prefer the canonical Converse fields, including explicit zero values.
  • Preserve the original response body and raw usage.

The prompt-token calculation follows the AWS prompt caching documentation:

prompt_tokens = inputTokens + cacheReadInputTokens + cacheWriteInputTokens

LiteLLM uses the same calculation in AmazonConverseConfig.transform_usage, lines 2339–2344:

input_tokens: Final = raw_input_tokens + cache_read_input_tokens + cache_creation_input_tokens
total_tokens: Final = usage.get("totalTokens", input_tokens + output_tokens)

It then assigns input_tokens to Usage.prompt_tokens. Alias precedence follows its separate cache-count helpers.

Tests

  • Add 12 usage cases covering cache reads, cache writes, aliases, zero and null values, and explicit totals in both response modes.
  • Add end-to-end assertions for the prompt and cache logging variables, using the existing Bedrock mock and selectable fixtures.
  • Compare the five normalized usage fields against LiteLLM commit 3d526bfe: all 120 comparisons across 12 cases and two response modes match.

Validation recorded when preparing this change: all 65 assertions in t/plugin/ai-proxy-bedrock-single.t passed, and luacheck reported no warnings or errors for the changed Lua code.

Checklist

  • I have explained the need for this PR and the problem it solves
  • I have explained the changes or the new features added to this PR
  • I have added tests corresponding to this change
  • I have updated the documentation to reflect this change (existing documentation already defines prompt tokens as including cache tokens; no new configuration or API)
  • I have verified that this change is backward compatible

@membphis membphis left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants